# -*- coding: utf-8 -*-
from functools import reduce
import os
import arcpy
import argparse
import json
import sys
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from utils.utools import str_to_bool
import pandas as pd
from datetime import datetime

def check_field_len(new_field, keep_head):
    mod_field = new_field
    if len(new_field) > 31:
        if keep_head:
            mod_field = new_field[:31]  # 取前31个字符
        else:
            mod_field = new_field[-31:] # 取后31个字符
        arcpy.AddMessage("字段名过长，已截断为：{}".format(mod_field))
    return mod_field

def join_df_func(main_df, right_df, how="left"):
    return main_df.join(right_df, how=how)

def safe_drop(df, cols_dropped, axis=1):
    exist_cols = [c for c in cols_dropped if c in df.columns]
    return df.drop(exist_cols, axis=axis) if exist_cols else df

def staType_to_staColName(staType):
    '''
        根据统计类型字符串返回统计列名。例如：MAXIMUM-->MAX，当统计类型是MAXIMUM时，ZonalStatisticsAsTables生成的列名是：MAX
    '''
    if not isinstance(staType, basestring):  # Python 2 中检查是否为字符串或unicode
        raise ValueError("Input must be a string")

    staType = staType.strip().upper()
    mapping = {
        'MEAN': 'MEAN',
        'MAXIMUM': 'MAX',
        'MINIMUM': 'MIN',
        'RANGE': 'RANGE',
        'STD': 'STD',
        'SUM': 'SUM',
    }
    return mapping.get(staType)  # 如果staType不在字典中，返回None


def pandas_to_feature_class(input_fc, df, output_fc, join_id_column="OBJECTID"):
    """
    将一个 Pandas DataFrame 的属性表附加到一个面要素类的几何上，并输出为新的要素类。
    (此版本已修正，支持连接键为DataFrame的列或索引，并兼容旧版Pandas和系统字段)
    """
    try:
        arcpy.AddMessage("正在复制要素几何: {} -> {}".format(input_fc, output_fc))
        arcpy.Copy_management(input_fc, output_fc)

        is_in_columns = join_id_column in df.columns
        is_index = df.index.name == join_id_column

        if not is_in_columns and not is_index:
            raise ValueError("DataFrame 中找不到名为 '{}' 的列或索引。".format(join_id_column))

        if is_in_columns:
            df_dict = df.set_index(join_id_column).to_dict('index')
            arcpy.AddMessage("已将列 '{}' 设置为索引并转换为字典。".format(join_id_column))
        else:
            df_dict = df.to_dict('index')
            arcpy.AddMessage("已将索引 '{}' 转换为字典。".format(join_id_column))

        arcpy.AddMessage("正在删除旧的属性字段...")
        fields_to_keep = ["OBJECTID", "Shape", "Shape_Length", "Shape_Area"]
        if join_id_column.upper() not in ["OBJECTID", "Shape_Length", "Shape_Area"]:
            fields_to_keep.append(join_id_column)

        fields = arcpy.ListFields(output_fc)
        for field in fields:
            if field.name not in fields_to_keep:
                arcpy.DeleteField_management(output_fc, field.name)

        arcpy.AddMessage("正在添加新的属性字段...")
        type_mapping = {
            'object': 'TEXT', 'int64': 'LONG', 'int32': 'LONG',
            'float64': 'DOUBLE', 'float32': 'FLOAT', 'bool': 'SHORT', 'datetime64[ns]': 'DATE'
        }

        for col_name in df.columns:
            if col_name.upper() in ["OBJECTID", "SHAPE"]:
                continue  # OBJECTID 和 Shape 是系统字段，不能添加

            pandas_type = str(df[col_name].dtype)
            arcpy_type = type_mapping.get(pandas_type, 'TEXT')
            field_length = 255
            if arcpy_type == 'TEXT':
                max_len = df[col_name].astype(str).str.len().max()
                if pd.notnull(max_len):
                    field_length = min(int(max_len * 1.5), 254)
            arcpy.AddField_management(output_fc, col_name, arcpy_type, field_length=field_length)

        arcpy.AddMessage("正在更新要素属性...")
        fields_to_update = [f.name for f in arcpy.ListFields(output_fc) if f.name not in ["OBJECTID", "Shape"]]

        cursor_fields = ["OBJECTID"] + fields_to_update

        with arcpy.da.UpdateCursor(output_fc, cursor_fields) as cursor:
            for row in cursor:
                oid = row[0]
                if oid in df_dict:
                    new_data = df_dict[oid]
                    for i, field_name in enumerate(fields_to_update):
                        row[i + 1] = new_data.get(field_name)
                    cursor.updateRow(row)

        arcpy.AddMessage("成功创建新要素类: {}".format(output_fc))
        return output_fc

    except Exception as e:
        arcpy.AddError("处理过程中发生错误: {}".format(e))
        if arcpy.Exists(output_fc):
            try:
                arcpy.Delete_management(output_fc)
                arcpy.AddWarning("已删除不完整的输出文件: {}".format(output_fc))
            except:
                pass
        raise

def Exact_RasterFactors_toSlopeUnits(input_slopeUnits_fc, rasters_statype_lst, slope_zonal_ID="OBJECTID"):
    arcpy.env.overwriteOutput = True

    in_fc_dir = os.path.dirname(input_slopeUnits_fc)
    arcpy.env.workspace = in_fc_dir
    stat_table_dfs = []
    out_table_list = []
    # 循环处理每个栅格
    for raster_path, statistics_type in rasters_statype_lst:
        # 获取栅格文件名（用于命名输出表格，避免重复）
        raster_name = arcpy.Describe(raster_path).baseName  # 如 "slope"、"dem"
        raster_name = raster_name[:-4] if raster_name.endswith(('.tif', '.jpg')) else raster_name   # 去掉后缀
        now_str = datetime.now().strftime('%d%H%M')
        out_table = "{}\\{}_{}stats".format(in_fc_dir, raster_name, now_str)  # 输出表格：slope_stats、dem_stats...
        arcpy.AddMessage("\nout_table: {}".format(out_table))
        arcpy.AddMessage("raster_name: {}".format(raster_name))
        arcpy.AddMessage("Step 1: 开始统计每个栅格图层......")
        try:
            # 处理中文路径（Python 2.7 必需）
            if isinstance(raster_path, unicode):
                raster_path = raster_path.encode("gbk")
            # 确保临时表不存在（避免冲突）
            if arcpy.Exists(out_table):
                arcpy.Delete_management(out_table)
            # 调用 ZonalStatisticsAsTable 处理单个栅格
            arcpy.sa.ZonalStatisticsAsTable(
                in_zone_data=input_slopeUnits_fc,
                zone_field=slope_zonal_ID,
                in_value_raster=raster_path,
                out_table=out_table,
                ignore_nodata='DATA',
                statistics_type=statistics_type
            )
            out_table_list.append(out_table)
            arcpy.AddMessage("Step 1.1 已完成栅格 {} 的统计，输出表格：{}".format(raster_name, out_table))

            arcpy.AddMessage("Step 1.2 使用pandas读取表格数据....")
            # （1）构建统计表的DataFrame
            field_names = [i.name for i in arcpy.ListFields(out_table)]
            cursor = arcpy.da.SearchCursor(out_table, field_names)
            df = pd.DataFrame(data=[row for row in cursor], columns=field_names)

            # （2）重命名统计列的名称。修正ID字段名（与主表一致）
            # df = df.rename(columns={c: slope_zonal_ID for c in df.columns if slope_zonal_ID in c})
            # 为统计字段添加栅格名前缀（如MEAN→predRe2024_MEAN，避免多栅格字段重名）
            statistics_type_lst = [statistics_type]
            for stat_type in statistics_type_lst:
                old_statType_colName = staType_to_staColName(stat_type)
                new_stat_field_name = "{}_{}".format(raster_name, old_statType_colName) # 旧的统计字段添加栅格名前缀（如MEAN→predRe2024_MEAN，避免多栅格字段重名）
                new_stat_field_name = check_field_len(new_stat_field_name, keep_head=False) # 避免字段名称过长，只截取后面的31个字符
                df = df.rename(columns={old_statType_colName: new_stat_field_name})
            # （3）删除一些不必要的列
            df = df.set_index(slope_zonal_ID, drop=True)
            df = safe_drop(df, ['OBJECTID_1', 'COUNT', 'AREA', 'Shape', 'Shape_Length', 'Shape_Area'])
            stat_table_dfs.append(df)

        except arcpy.ExecuteError as e:
            error_msg = e.message
            arcpy.AddMessage("处理栅格 {} 失败：{}".format(raster_name, error_msg))
        except Exception as e:
            # 捕获其他非ArcPy错误
            arcpy.AddError("程序错误：\n{}".format(e))

    # 4. 读取斜坡单元属性表为DataFrame
    arcpy.AddMessage("Step 2: 正在读取斜坡单元要素类至DataFrame:{}".format(input_slopeUnits_fc))
    slopeUnits_field_names = [i.name for i in arcpy.ListFields(input_slopeUnits_fc)]
    # slopeUnits_field_names.append("SHAPE@XY")
    data = [row for row in arcpy.da.SearchCursor(input_slopeUnits_fc, slopeUnits_field_names)]
    df_slopeUnits = pd.DataFrame(data, columns=slopeUnits_field_names)
    df_slopeUnits = df_slopeUnits.set_index(slope_zonal_ID, drop=True)


    if len(stat_table_dfs) > 0:
        arcpy.AddMessage("Step 3: 正在合并斜坡单元+所有统计表，并输出为要素类文件......")
        # 主表 + 所有统计表
        result_df = reduce(join_df_func, [df_slopeUnits] + stat_table_dfs)
        # arcpy.AddMessage('\n\t====result_df====')
        # arcpy.AddMessage(result_df.head(2))
        in_fc_name = os.path.basename(input_slopeUnits_fc)

        now_str = datetime.now().strftime('%d%H%M')
        out_fc_name = "{}_{}copyed".format(in_fc_name, now_str)
        out_fc_fullpath = os.path.join(in_fc_dir, out_fc_name)
        pandas_to_feature_class(input_slopeUnits_fc, result_df, out_fc_fullpath, join_id_column=slope_zonal_ID)

    arcpy.AddMessage("Step 4: 正在清理：删除临时的统计表......")
    for out_table in out_table_list:
        if arcpy.Exists(out_table):
            arcpy.Delete_management(out_table)
    arcpy.AddMessage("最终输出结果: 【{}】".format(out_fc_fullpath))

def main():
    """
    主函数，用于解析参数并执行逻辑。
    """
    # 1. 设置参数解析器
    parser = argparse.ArgumentParser(description="从栅格提取数据到斜坡单元。")
    parser.add_argument(
        "--fc_slopeunits",
        required=True,
        help="斜坡单元要素类的路径。"
    )
    parser.add_argument(
        "--stat_json",
        required=True,
        help="一个JSON格式的字符串，包含栅格路径和统计类型。"
    )

    parser.add_argument('--id_field', required=True, help='斜坡单元的ID字段')
    parser.add_argument('--del_existing_fields', required=True, help='是否删除已有字段')

    # 2. 解析参数
    # 此时，args.fc_slopeunits 和 args.stat_json 都是 str 类型 (mbcs 编码的字节串)
    args = parser.parse_args()

    # arcpy.AddMessage("type(args.stat_json):{}".format(type(args.stat_json)))    # type(args.stat_json):<type 'str'>
    # arcpy.AddMessage("type(args.fc_slopeunits):{}".format(type(args.fc_slopeunits))) # type(args.fc_slopeunits):<type 'str'>

    stat_list = json.loads(args.stat_json, encoding='mbcs')
    # arcpy.AddMessage("成功解析JSON列表，类型为: {}, 内容为: {}".format(type(stat_list), stat_list))

    statics_settings = []
    for sta in stat_list:
        # arcpy.AddMessage("{},{}".format(type(sta[0]), type(sta[1])))    # <type 'unicode'>,<type 'unicode'>
        # arcpy.AddMessage("{},{}".format(sta[0].encode('mbcs'), sta[1].encode('mbcs')))
        statics_settings.append((sta[0].encode('mbcs'), sta[1].encode('mbcs')))
    # arcpy.AddMessage(statics_settings)
    del_existing_fields = str_to_bool(args.del_existing_fields)
    # arcpy.AddMessage("args.id_field: {}".format(args.id_field))
    # arcpy.AddMessage("del_existing_fields: {}".format(del_existing_fields))
    # exact_factors_forSlopeUnits(args.fc_slopeunits, slope_zonal_ID=args.id_field, del_exist_fields=del_existing_fields, statics_settings=statics_settings)
    Exact_RasterFactors_toSlopeUnits(args.fc_slopeunits, statics_settings, slope_zonal_ID=args.id_field)

def tt():
    # 工具箱模式
    input_fc = arcpy.GetParameterAsText(0)
    slope_zonal_ID = arcpy.GetParameterAsText(1)
    del_exist_fields = arcpy.GetParameter(2)

    current_file_dir = os.path.split(os.path.realpath(__file__))[0]
    results_dir = current_file_dir + "/settings"
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    # # Set the current workspace
    # arcpy.env.workspace = "G:/DataForDoctorPaper/博士论文数据.gdb"
    # input_fc = 'SlopeUnits_V5'
    # slope_zonal_ID = 'OBJECTID'
    # year = 2003
    # del_exist_fields = True
    # data_version = "V5"
    # results_dir = "./settings"

    statics_settings = "{0}/slopeSettings_{1}.csv".format(results_dir)
    exact_factors_forSlopeUnits(input_fc=input_fc,
                                slope_zonal_ID=slope_zonal_ID,
                                del_exist_fields=del_exist_fields,
                                statics_settings=statics_settings)

if __name__=="__main__":
    main()
